#!/usr/bin/env python
# coding: utf-8

# In[8]:


import pandas as pd
import numpy as np
import tushare as ts
from datetime import datetime, timedelta

# 设置tushare token（需要先在tushare官网注册获取）
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()

# 获取上一个交易日
def get_last_trade_date():
    today = datetime.now()
    # 获取最近一个交易日
    trade_date = pro.trade_cal(exchange='', start_date=(today - timedelta(days=7)).strftime('%Y%m%d'), 
                             end_date=today.strftime('%Y%m%d'), is_open='1')
    last_trade_date = trade_date.iloc[-1]['cal_date']
    return last_trade_date

last_trade_date = get_last_trade_date()
print(f"上一个交易日: {last_trade_date}")


# In[9]:


# 获取所有股票的基本信息
stock_list = pro.stock_basic(exchange='', list_status='L')
stock_codes = stock_list['ts_code'].tolist()

# 获取上一个交易日的日线数据
def get_daily_data(trade_date):
    # 使用tushare的daily接口获取数据
    daily_data = pro.daily(trade_date=trade_date)
    return daily_data

daily_df = get_daily_data(last_trade_date)

# 合并股票基本信息
merged_df = pd.merge(daily_df, stock_list, on='ts_code', how='left')


# In[10]:


# 数据概览
print("数据概览:")
print(f"股票数量: {len(merged_df)}")
print(f"数据时间: {last_trade_date}")
print("\n前5条数据:")
print(merged_df.head())

# 基本统计信息
print("\n基本统计信息:")
print(merged_df.describe())

# 各行业股票数量统计（如果数据中包含行业信息）
if 'industry' in merged_df.columns:
    industry_count = merged_df['industry'].value_counts()
    print("\n各行业股票数量:")
    print(industry_count)

# 涨跌幅分布分析
print("\n涨跌幅分布:")
print(merged_df['pct_chg'].describe())

# 成交量分析
print("\n成交量(手)分析:")
print(merged_df['vol'].describe())

# 成交额分析
print("\n成交额(千元)分析:")
print(merged_df['amount'].describe())


# In[18]:


import pandas as pd
import numpy as np

# 假设我们已经有了merged_df，包含上一个交易日的股票数据
# 首先进行数据清洗

def preprocess_data(df):
    # 删除缺失值
    df = df.dropna()
    
    # 转换日期格式
    df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
    
    # 按股票代码和日期排序，为计算技术指标做准备
    df = df.sort_values(['ts_code', 'trade_date'])
    
    # 重置索引
    df = df.reset_index(drop=True)
    
    return df

# 预处理数据
processed_df = preprocess_data(merged_df.copy())
print("预处理后的数据前5行:")
print(processed_df.head())

def calculate_ma(df, window=5):
    """
    计算移动平均线
    :param df: 包含股票数据的DataFrame
    :param window: 移动平均窗口大小
    :return: 添加了MA列的DataFrame
    """
    # 按股票代码分组计算
    df[f'ma{window}'] = df.groupby('ts_code')['close'].transform(
        lambda x: x.rolling(window=window, min_periods=1).mean()
    )
    return df

# 计算MA5和MA10
processed_df = calculate_ma(processed_df, 5)
processed_df = calculate_ma(processed_df, 10)

# 可选：计算更多移动平均线
processed_df = calculate_ma(processed_df, 20)  # MA20
processed_df = calculate_ma(processed_df, 60)  # MA60


def calculate_rsi(df, window=14):
    """
    计算相对强弱指数(RSI)
    :param df: 包含股票数据的DataFrame
    :param window: RSI计算窗口
    :return: 添加了RSI列的DataFrame
    """
    # 按股票代码分组
    grouped = df.groupby('ts_code')['close']
    
    # 计算价格变化
    delta = grouped.diff()
    
    # 分离上涨和下跌
    up = delta.clip(lower=0)
    down = -1 * delta.clip(upper=0)
    
    # 计算平均上涨和下跌
    roll_up = up.transform(lambda x: x.rolling(window).mean())
    roll_down = down.transform(lambda x: x.rolling(window).mean())
    
    # 计算相对强度
    rs = roll_up / roll_down
    
    # 计算RSI
    df[f'rsi{window}'] = 100 - (100 / (1 + rs))
    
    return df

# 计算14日RSI
processed_df = calculate_rsi(processed_df, 14)

# 可选：计算其他周期的RSI
processed_df = calculate_rsi(processed_df, 6)   # RSI6
processed_df = calculate_rsi(processed_df, 24)  # RSI24


def calculate_macd(df, fast=12, slow=26, signal=9):
    """
    计算MACD指标
    :param df: 包含股票数据的DataFrame
    :param fast: 快速EMA周期
    :param slow: 慢速EMA周期
    :param signal: 信号线周期
    :return: 添加了MACD相关列的DataFrame
    """
    # 按股票代码分组
    grouped = df.groupby('ts_code')['close']
    
    # 计算EMA
    ema_fast = grouped.transform(lambda x: x.ewm(span=fast, adjust=False).mean())
    ema_slow = grouped.transform(lambda x: x.ewm(span=slow, adjust=False).mean())
    
    # 计算DIF
    df['macd_dif'] = ema_fast - ema_slow
    
    # 计算DEA
    df['macd_dea'] = df.groupby('ts_code')['macd_dif'].transform(
        lambda x: x.ewm(span=signal, adjust=False).mean()
    )
    
    # 计算MACD柱状图
    df['macd'] = 2 * (df['macd_dif'] - df['macd_dea'])
    
    return df

# 计算MACD
processed_df = calculate_macd(processed_df)


# In[ ]:




